import paramiko
import os

# ssh = paramiko.SSHClient()
# ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# ssh.connect(hostname="10.36.176.168",
#             username="root",
#             password="1")
# # 执行完成后，返回3个结果，正确输出、错误输出、标准输入
# _stdin, stdout, stderr = ssh.exec_command("ip a")
# stdout = stdout.read().decode()
# print(stdout if stdout else stderr.read().decode())


# 封装transport
def remote_modify(mode="ssh", host="127.0.0.1", port=22, user="root", password=None):
    # 完成远程连接
    transport = paramiko.Transport((host, port))
    transport.connect(username=user, password=password)
    if mode == "ssh":
        # 远程连接操作
        ssh = paramiko.SSHClient()
        ssh._transport = transport  # 将ssh连接和上面的账号密码等信息绑定
        # 模仿终端
        while True:
            # 获取主机名
            _, hostname, _stderr1 = ssh.exec_command("hostname")
            tag = "#" if user == "root" else "$"
            _, path, _stderr2 = ssh.exec_command("pwd")
            print(f"[{user}@{hostname.read().decode().split('.')[0]} {os.path.basename(path.read().decode())}]{tag} ".replace("\n", ""), end="")
            # print("输入quit，退出终端")
            cmd = input(
            )
            a, b, c = ssh.exec_command(cmd)
            stdout = b.read().decode()
            print(stdout if stdout else c.read().decode())
    elif mode == "sftp":
        # 远程上传下载
        sftp = paramiko.SFTPClient.from_transport(transport)
        r = input("请问你是要上传还是下载：")
        if r == "上传":
            local, remote = input("请输入要上传的文件路径："), input("请输入上传的位置：")
            sftp.put(localpath=local, remotepath=remote)
            print(f"{local} ========上传成功========> {remote}")
        else:
            remote, local = input("请输入要下载的文件路径："), input("请输入保存的位置：")
            sftp.get(remotepath=remote, localpath=local)
            print(f"{remote} ========下载成功========> {local}")
        

remote_modify(mode="sftp", host="10.36.176.168", password="1")
